from openai import OpenAI

from utils.llm_setting import Api_Key, Api_Url, Model_List


class LLMApi:
    def __init__(self):
        self.api_key: str = Api_Key
        self.api_url: str = Api_Url
        self.model_list: list = Model_List
        pass

    def get_api_key(self) -> str:
        return self.api_key

    def get_api_url(self) -> str:
        return self.api_url

    def get_model_list(self) -> list:
        return self.model_list

    def set_api_key(self, api_key: str) -> None:
        self.api_key = api_key

    def set_api_url(self, api_url: str) -> None:
        self.api_url = api_url

    def set_model_list(self, model_list: list) -> None:
        self.model_list = model_list

    def get_model_res(self, prompt: str) -> dict:
        
        client = OpenAI(api_key=self.api_key, base_url=self.api_url)

        # 发送非流式输出的请求
        messages = [{"role": "user", "content": prompt}]
        response = client.chat.completions.create(
            model=self.model_list[0],  # "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
            messages=messages,
            stream=False,
            max_tokens=4096,
        )
        content = response.choices[0].message.content
        reasoning_content = response.choices[0].message.reasoning_content

        # Round 2
        messages.append({"role": "assistant", "content": content})
        messages.append({"role": "user", "content": "继续"})
        response = client.chat.completions.create(
            model=self.model_list[0], messages=messages, stream=False
        )

        second_round_content = response.choices[0].message.content

        # 合并结果并过滤掉思考过程
        full_content = (reasoning_content or "") + (second_round_content or "")
        
        # 去除"思考："、"推理："等前缀的思考过程内容
        import re
        filtered_content = re.sub(r'(思考|推理|Thought|Reasoning)[：:].*?\n', '', full_content)
        filtered_content = re.sub(r'\[.*?\]', '', filtered_content)  # 去除方括号内容
        
        return filtered_content.strip()
